±1 Operation 2

/icons/hr.icon

問題

/icons/hr.icon

出典

ABC255 - D 788 diff

/icons/hr.icon

考察

配列$ A はソートしても一般性を失わないことはある程度競技プログラミングをやって来ればすぐにわかる

したがって、以下 $ Aはソートして昇順で扱う

次に、各クエリでの$ X_iがソートした$ Aのどこに来るのかを考える

つまり、$ A_1 \leq A_2 \leq ... \leq A_K \leq X_i \leq A_{K + 1} \leq ... \leq A_N を満たす$ Kを探索する

これは Python であれば bisectモジュールを使用すれば$ O(\log N)すぐに求められる

この場合に必要な操作の回数を考える

すると、$ A_1 〜 A_Kまでは $ X_i以下であり、$ A_{K + 1} 〜 A_Nは$ X_i以上であることがわかるので以下の式で操作回数は求められる

$ \{ (X_i - A_1) + (X_i - A_2) + ... + (X_i - A_K) \} + \{ (A_{K + 1} - X_i) + (A_{K + 2} - X_i) + ... + (A_N - X_i) \}

$ = \{ K \cdot X_i - (A_1 + A_2 + ... + A_K) \} + \{ (A_{K + 1} + A_{K + 2} + ... + A_N) - (N - K) \cdot X_i \} \\

$ = K \cdot X_i - \sum_{j = 1}^{K} A_j + \sum_{j = K + 1}^{N} A_j - (N - K) \cdot X_i

$ = \sum_{j = K + 1}^{N} A_j - \sum_{j = 1}^{K} A_j + (2K - N) \cdot X_i

$ = \left( \sum_{j = 1}^{N} A_j - \sum_{j = 1}^{K} A_j \right) - \sum_{j = 1}^{K} A_j + (2K - N) \cdot X_i

$ = \left( \sum_{j = 1}^{N} A_j - 2 \sum_{j = 1}^{K} A_j \right) + (2K - N) \cdot X_i

この$ \sum_{j = 1}^{K} A_jは$ Aに対する累積和 $ Sであらかじめ前処理で求めておくと、各クエリに関して$ O(\log N)で求めることができる

したがって、各クエリは $ O(\log N)で求めることが出来たので、計算量は $ O \left((N + Q) \log N \right)となりました

/icons/hr.icon

コード

code:Python

from bisect import bisect_right

n, q = map(int, input().split())

a = sorted(list(map(int, input().split())))

s = 0 * (n + 1)

for i in range(n):

si + 1 = ai + si

sum_ = sum(a)

for _ in range(q):

X = int(input())

K = bisect_right(a, X) # O(log N)

print(sum_ - 2 * sK + (2 * K - n) * X)